package datastructure.tree.bst2;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * 二叉搜索树-范围查找
 *
 * @author 杨 强
 * @package datastructure.tree.bst2
 * @createTime 2023-03-21 21:50
 */
public class BST01Scope {
    public static void main(String[] args) {
        /*
                        3
                       / \
                      2   4
                     /
                    1
         */
        TreeNode node1 = new TreeNode(1, "1", null, null);
        TreeNode node2 = new TreeNode(2, "2", node1, null);
        TreeNode node4 = new TreeNode(4, "4", null, null);
        TreeNode node3 = new TreeNode(3, "3", node2, node4);
        BST01Scope bst01 = new BST01Scope(node3);

        System.out.println(bst01.less(4));
        System.out.println(bst01.greater(1));
        System.out.println(bst01.greater1(1));
        System.out.println(bst01.between(1, 3));

    }

    /**
     * 二叉搜索树的根节点
     */
    TreeNode root;

    /**
     * 构造函数-生成二叉树
     * @param root
     */
    public BST01Scope(TreeNode root) {
        this.root = root;
    }

    /**
     * 找比key小的key的元素的值
     * @param key key
     * @return
     */
    public List<Object> less(int key) {
        List<Object> result = new ArrayList<>();

        // 先进行中序遍历
        // 然后进行和目标key进行比较, 比key小的就放入集合, 遍历到比key大的就退出
        // 创建一个栈来进行辅助遍历
        LinkedList<TreeNode> stack = new LinkedList<>();
        TreeNode p = root;
        while (p != null || !stack.isEmpty()){
            // 当前树的左孩子不为空就一直找,直到把这个子树找到底
            if (p != null){
                // 入栈
                stack.push(p);
                // 往左找
                p = p.left;
            }else {
                // 左子树找完了就找右子树
                // 出栈
                TreeNode pop = stack.pop();
                if (pop.key < key) {
                    result.add(pop.value);
                }else {
                    // 出现比key大就结束
                    break;
                }
                // 上面的代码就处理完了左孩子和自己本身, 然后开始处理右孩子
                p = pop.right;
            }
        }

        return result;
    }

    /**
     * 找比key大的key的元素的值-中序遍历
     * @param key key
     * @return
     */
    public List<Object> greater(int key) {
        List<Object> result = new ArrayList<>();

        LinkedList<TreeNode> stack = new LinkedList<>();
        TreeNode p = root;
        while (p != null || !stack.isEmpty()){
            // 当前树的左孩子不为空就一直找,直到把这个子树找到底
            if (p != null){
                // 入栈
                stack.push(p);
                // 往左找
                p = p.left;
            }else {
                // 左子树找完了就找右子树
                // 出栈
                TreeNode pop = stack.pop();
                // 这样的方式的缺点就是比自己小的明明是不要,但是也遍历完了
                if (pop.key > key) {
                    result.add(pop.value);
                }
                // 上面的代码就处理完了左孩子和自己本身, 然后开始处理右孩子
                p = pop.right;
            }
        }
        return result;
    }

    /**
     * 找比key大的key的元素的值-使用反向中序遍历优化
     * @param key key
     * @return
     */
    public List<Object> greater1(int key) {

        List<Object> result = new ArrayList<>();

        LinkedList<TreeNode> stack = new LinkedList<>();
        TreeNode p = root;
        while (p != null || !stack.isEmpty()){
            // 当前树的右孩子不为空就一直找,直到把这个子树找到底
            if (p != null){
                // 入栈
                stack.push(p);
                // 往右找
                p = p.right;
            }else {
                // 右子树找完了就处理找右子树
                // 出栈
                TreeNode pop = stack.pop();
                if (pop.key > key) {
                    result.add(pop.value);
                }else {
                    // 出现比自己小的就退出
                    break;
                }
                // 上面的代码就处理完了右孩子和自己本身, 然后开始处理左孩子
                p = pop.left;
            }
        }
        return result;

    }

    /**
     * 找大于等于key1, 小于等于key2的key的元素值
     * @param key1 key
     * @param key2 key
     * @return
     */
    public List<Object> between(int key1, int key2) {
        List<Object> result = new ArrayList<>();

        LinkedList<TreeNode> stack = new LinkedList<>();
        TreeNode p = root;
        while (p != null || !stack.isEmpty()){
            // 当前树的右孩子不为空就一直找,直到把这个子树找到底
            if (p != null){
                // 入栈
                stack.push(p);
                // 往左找
                p = p.left;
            }else {
                // 左子树找完了就处理找右子树
                // 出栈
                TreeNode pop = stack.pop();
                if (pop.key >= key1 && pop.key <= key2) {
                    result.add(pop.value);
                }else if (pop.key > key2){
                    // 不满足范围了直接退出
                    break;
                }
                // 上面的代码就处理完了左孩子和自己本身, 然后开始处理右孩子
                p = pop.right;
            }
        }

        return result;
    }
}
